package main

import (
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

// 下载任务结构
type DownloadTask struct {
	URL          string
	TargetPath   string
	TotalSize    int64
	Downloaded   int64
	SegmentCount int
	Segments     []SegmentInfo
	StartTime    time.Time
	Speed        float64
	Completed    bool
	Error        error
	Progress     float64
	SpeedLimit   int64
	Mutex        sync.Mutex
}

// 分片信息
type SegmentInfo struct {
	Index      int
	StartByte  int64
	EndByte    int64
	Downloaded int64
	Completed  bool
	TempFile   string
}

// 限速读取器
type ThrottledReader struct {
	reader        io.Reader
	bytesPerSec   int64
	bytesRead     int64
	lastCheckTime time.Time
}

// 创建新的限速读取器
func NewThrottledReader(reader io.Reader, bytesPerSec int64) *ThrottledReader {
	return &ThrottledReader{
		reader:        reader,
		bytesPerSec:   bytesPerSec,
		lastCheckTime: time.Now(),
	}
}

// 实现 Read 接口
func (t *ThrottledReader) Read(p []byte) (n int, err error) {
	n, err = t.reader.Read(p)
	if n > 0 && t.bytesPerSec > 0 {
		t.bytesRead += int64(n)

		// 检查是否需要限速
		elapsed := time.Since(t.lastCheckTime)
		if elapsed.Seconds() >= 1.0 {
			// 重置计数器
			t.bytesRead = 0
			t.lastCheckTime = time.Now()
		} else if t.bytesRead > t.bytesPerSec {
			// 计算需要等待的时间
			expectedTime := float64(t.bytesRead) / float64(t.bytesPerSec)
			waitTime := time.Duration((expectedTime - elapsed.Seconds()) * float64(time.Second))
			if waitTime > 0 {
				time.Sleep(waitTime)
			}
			// 重置计数器
			t.bytesRead = 0
			t.lastCheckTime = time.Now()
		}
	}
	return n, err
}

// 获取文件信息
func getFileInfo(url string) (int64, bool, error) {
	// 发送 HEAD 请求获取文件信息
	req, err := http.NewRequest("HEAD", url, nil)
	if err != nil {
		return 0, false, err
	}
	// 设置 User-Agent 避免被某些服务器拒绝
	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36")

	client := &http.Client{
		Timeout: 30 * time.Second,
	}
	resp, err := client.Do(req)
	if err != nil {
		return 0, false, err
	}
	defer resp.Body.Close()
	// 检查状态码
	if resp.StatusCode != http.StatusOK {
		return 0, false, fmt.Errorf("服务器返回错误状态码: %d", resp.StatusCode)
	}

	// 获取文件大小
	contentLength := resp.ContentLength

	// 检查是否支持断点续传
	acceptRanges := resp.Header.Get("Accept-Ranges")
	supportsRanges := acceptRanges == "bytes"

	return contentLength, supportsRanges, nil
}

// 创建下载任务
func createDownloadTask(url, targetPath string, threads int) (*DownloadTask, error) {
	// 获取文件信息
	fileSize, supportsRanges, err := getFileInfo(url)
	if err != nil {
		return nil, err
	}

	// 如果不支持范围请求但用户要求分段下载，则回退到单线程下载
	segmented := threads > 1
	if segmented && !supportsRanges {
		segmented = false
		threads = 1
		fmt.Println("警告: 服务器不支持断点续传，将使用单线程下载")
	}

	// 创建新任务
	task := &DownloadTask{
		URL:        url,
		TargetPath: targetPath,
		TotalSize:  fileSize,
		StartTime:  time.Now(),
	}

	// 如果支持分段下载，则创建分段信息
	if segmented {
		// 计算分段大小和数量
		segmentCount := threads
		segmentSize := fileSize / int64(segmentCount)

		// 确保每个分段至少有 1MB
		if segmentSize < 1024*1024 && fileSize > 1024*1024 {
			segmentCount = int(fileSize / (1024 * 1024))
			if segmentCount < 1 {
				segmentCount = 1
			}
			segmentSize = fileSize / int64(segmentCount)
		}

		task.SegmentCount = segmentCount
		task.Segments = make([]SegmentInfo, segmentCount)

		// 初始化每个分段的信息
		for i := 0; i < segmentCount; i++ {
			startByte := int64(i) * segmentSize
			endByte := startByte + segmentSize - 1

			// 最后一个分段可能会更大一些
			if i == segmentCount-1 {
				endByte = fileSize - 1
			}

			task.Segments[i] = SegmentInfo{
				Index:     i,
				StartByte: startByte,
				EndByte:   endByte,
				TempFile:  fmt.Sprintf("%s.part%d", targetPath, i),
			}
		}
	}

	return task, nil
}

// 下载单个分段
func (task *DownloadTask) downloadSegment(segment *SegmentInfo, wg *sync.WaitGroup, errorChan chan error) {
	defer wg.Done()

	// 创建请求
	req, err := http.NewRequest("GET", task.URL, nil)
	if err != nil {
		errorChan <- err
		return
	}

	// 设置范围头和 User-Agent
	rangeHeader := fmt.Sprintf("bytes=%d-%d", segment.StartByte, segment.EndByte)
	req.Header.Set("Range", rangeHeader)
	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36")

	// 发送请求
	client := &http.Client{
		Timeout: 0, // 无超时
	}
	resp, err := client.Do(req)
	if err != nil {
		errorChan <- err
		return
	}
	defer resp.Body.Close()

	// 检查状态码
	if resp.StatusCode != http.StatusPartialContent {
		errorChan <- fmt.Errorf("服务器返回错误状态码: %d", resp.StatusCode)
		return
	}

	// 创建临时文件
	tempFile, err := os.Create(segment.TempFile)
	if err != nil {
		errorChan <- err
		return
	}
	defer tempFile.Close()

	// 准备读取器，如果设置了速度限制则使用限速读取器
	var reader io.Reader = resp.Body
	if task.SpeedLimit > 0 {
		// 为每个线程分配速度限制
		perThreadLimit := task.SpeedLimit / int64(task.SegmentCount)
		if perThreadLimit < 1024 { // 确保至少有 1KB/s
			perThreadLimit = 1024
		}
		reader = NewThrottledReader(resp.Body, perThreadLimit)
	}

	// 读取响应并写入临时文件
	buffer := make([]byte, 32*1024) // 32KB 缓冲区
	for {
		n, err := reader.Read(buffer)
		if n > 0 {
			// 写入文件
			if _, writeErr := tempFile.Write(buffer[:n]); writeErr != nil {
				errorChan <- writeErr
				return
			}

			// 更新下载进度
			task.Mutex.Lock()
			segment.Downloaded += int64(n)
			task.Downloaded += int64(n)
			task.Progress = float64(task.Downloaded) / float64(task.TotalSize) * 100

			// 计算下载速度
			elapsed := time.Since(task.StartTime).Seconds()
			if elapsed > 0 {
				task.Speed = float64(task.Downloaded) / elapsed
			}
			task.Mutex.Unlock()
		}

		if err != nil {
			if err == io.EOF {
				break
			}
			errorChan <- err
			return
		}
	}

	// 标记分段完成
	task.Mutex.Lock()
	segment.Completed = true
	task.Mutex.Unlock()
}

// 单线程下载
func (task *DownloadTask) downloadSingle() error {
	// 创建请求
	req, err := http.NewRequest("GET", task.URL, nil)
	if err != nil {
		return err
	}

	// 设置 User-Agent
	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36")

	// 发送请求
	client := &http.Client{
		Timeout: 0, // 无超时
	}
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	// 检查状态码
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("服务器返回错误状态码: %d", resp.StatusCode)
	}

	// 创建目标文件
	err = os.MkdirAll(filepath.Dir(task.TargetPath), 0755)
	if err != nil {
		return err
	}

	file, err := os.Create(task.TargetPath)
	if err != nil {
		return err
	}
	defer file.Close()

	// 准备读取器，如果设置了速度限制则使用限速读取器
	var reader io.Reader = resp.Body
	if task.SpeedLimit > 0 {
		reader = NewThrottledReader(resp.Body, task.SpeedLimit)
	}

	// 创建缓冲区
	buffer := make([]byte, 32*1024) // 32KB 缓冲区

	// 读取响应并写入文件
	for {
		n, err := reader.Read(buffer)
		if n > 0 {
			// 写入文件
			if _, writeErr := file.Write(buffer[:n]); writeErr != nil {
				return writeErr
			}

			// 更新下载进度
			task.Mutex.Lock()
			task.Downloaded += int64(n)
			task.Progress = float64(task.Downloaded) / float64(task.TotalSize) * 100

			// 计算下载速度
			elapsed := time.Since(task.StartTime).Seconds()
			if elapsed > 0 {
				task.Speed = float64(task.Downloaded) / elapsed
			}
			task.Mutex.Unlock()
		}

		if err != nil {
			if err == io.EOF {
				break
			}
			return err
		}
	}

	return nil
}

// 合并分段文件
func (task *DownloadTask) mergeSegments() error {
	// 创建目标文件目录
	err := os.MkdirAll(filepath.Dir(task.TargetPath), 0755)
	if err != nil {
		return err
	}

	// 创建目标文件
	destFile, err := os.Create(task.TargetPath)
	if err != nil {
		return err
	}
	defer destFile.Close()

	// 按顺序合并所有分段
	for i := 0; i < task.SegmentCount; i++ {
		segment := &task.Segments[i]

		// 打开临时文件
		tempFile, err := os.Open(segment.TempFile)
		if err != nil {
			return err
		}

		// 复制内容到目标文件
		_, err = io.Copy(destFile, tempFile)
		tempFile.Close()

		// 删除临时文件
		os.Remove(segment.TempFile)

		if err != nil {
			return err
		}
	}

	return nil
}

// 检查是否可以恢复下载
func (task *DownloadTask) canResume() bool {
	if task.SegmentCount <= 0 {
		return false
	}

	// 检查每个分段的临时文件是否存在
	for i := 0; i < task.SegmentCount; i++ {
		segment := &task.Segments[i]
		info, err := os.Stat(segment.TempFile)
		if err != nil || info.Size() == 0 {
			return false
		}

		// 更新已下载的字节数
		segment.Downloaded = info.Size()
		task.Downloaded += info.Size()
	}

	return true
}

// 恢复下载
func (task *DownloadTask) resumeDownload() error {
	var wg sync.WaitGroup
	errorChan := make(chan error, task.SegmentCount)

	// 重新计算每个分段的起始位置
	for i := 0; i < task.SegmentCount; i++ {
		segment := &task.Segments[i]

		// 如果分段已完全下载，则跳过
		if segment.Downloaded >= (segment.EndByte - segment.StartByte + 1) {
			segment.Completed = true
			continue
		}

		// 调整起始字节位置
		segment.StartByte += segment.Downloaded

		wg.Add(1)
		go task.downloadSegment(segment, &wg, errorChan)
	}

	// 等待所有分段完成
	wg.Wait()
	close(errorChan)

	// 检查是否有错误
	for err := range errorChan {
		if err != nil {
			return err
		}
	}

	// 合并所有分段
	return task.mergeSegments()
}

// 开始下载
func (task *DownloadTask) startDownload() error {
	// 如果是分段下载
	if task.SegmentCount > 0 {
		var wg sync.WaitGroup
		errorChan := make(chan error, task.SegmentCount)

		// 启动所有分段的下载
		for i := 0; i < task.SegmentCount; i++ {
			wg.Add(1)
			go task.downloadSegment(&task.Segments[i], &wg, errorChan)
		}

		// 等待所有分段完成
		wg.Wait()
		close(errorChan)

		// 检查是否有错误
		for err := range errorChan {
			if err != nil {
				task.Error = err
				return err
			}
		}

		// 合并所有分段
		if err := task.mergeSegments(); err != nil {
			task.Error = err
			return err
		}
	} else {
		// 单线程下载
		if err := task.downloadSingle(); err != nil {
			task.Error = err
			return err
		}
	}

	// 标记任务完成
	task.Completed = true
	return nil
}

// 监控下载进度
func (task *DownloadTask) monitorProgress() {
	fmt.Println("下载开始...")

	// 清除之前的进度显示
	clearLine := func() {
		fmt.Print("\r                                                                               \r")
	}

	for !task.Completed && task.Error == nil {
		task.Mutex.Lock()
		progress := task.Progress
		speed := task.Speed / 1024 / 1024                    // 转换为 MB/s
		downloaded := float64(task.Downloaded) / 1024 / 1024 // 转换为 MB
		total := float64(task.TotalSize) / 1024 / 1024       // 转换为 MB
		task.Mutex.Unlock()

		// 计算剩余时间
		var eta string
		if speed > 0 {
			remainingBytes := float64(task.TotalSize - task.Downloaded)
			remainingSeconds := remainingBytes / (speed * 1024 * 1024)
			eta = formatTime(remainingSeconds)
		} else {
			eta = "计算中..."
		}

		// 打印进度条
		width := 40
		completed := int(progress / 100 * float64(width))

		clearLine()
		fmt.Printf("[")
		for i := 0; i < width; i++ {
			if i < completed {
				fmt.Print("=")
			} else if i == completed {
				fmt.Print(">")
			} else {
				fmt.Print(" ")
			}
		}
		fmt.Printf("] %.1f%% | %.2f/%.2f MB | %.2f MB/s | ETA: %s",
			progress, downloaded, total, speed, eta)

		time.Sleep(500 * time.Millisecond)
	}

	clearLine()

	// 打印最终状态
	if task.Error != nil {
		fmt.Printf("下载失败: %v\n", task.Error)
	} else {
		fmt.Printf("下载完成: %s\n", task.TargetPath)
		fmt.Printf("文件大小: %.2f MB\n", float64(task.TotalSize)/1024/1024)
		fmt.Printf("用时: %s\n", formatTime(time.Since(task.StartTime).Seconds()))
	}
}

// 格式化时间
func formatTime(seconds float64) string {
	hours := int(seconds) / 3600
	minutes := (int(seconds) % 3600) / 60
	secs := int(seconds) % 60

	if hours > 0 {
		return fmt.Sprintf("%d小时%d分%d秒", hours, minutes, secs)
	} else if minutes > 0 {
		return fmt.Sprintf("%d分%d秒", minutes, secs)
	}
	return fmt.Sprintf("%d秒", secs)
}

// 从URL中提取文件名
func getFilenameFromURL(url string) string {
	// 移除查询参数
	if idx := strings.Index(url, "?"); idx != -1 {
		url = url[:idx]
	}

	// 提取文件名
	parts := strings.Split(url, "/")
	filename := parts[len(parts)-1]

	// 如果文件名为空，使用默认名称
	if filename == "" {
		return "download.file"
	}

	return filename
}

func main() {
	// 解析命令行参数
	url := flag.String("url", "", "下载链接")
	output := flag.String("output", "", "保存路径")
	threads := flag.Int("threads", 4, "下载线程数")
	speedLimit := flag.Int64("limit", 0, "下载速度限制 (KB/s, 0表示不限制)")
	flag.Parse()

	// 检查必要参数
	if *url == "" {
		fmt.Println("请提供下载链接 (-url)")
		flag.Usage()
		return
	}

	// 如果没有指定输出文件名，从URL中提取
	if *output == "" {
		*output = getFilenameFromURL(*url)
	}

	// 创建下载任务
	task, err := createDownloadTask(*url, *output, *threads)
	if err != nil {
		fmt.Printf("创建下载任务失败: %v\n", err)
		return
	}

	// 设置速度限制
	if *speedLimit > 0 {
		task.SpeedLimit = *speedLimit * 1024 // 转换为字节/秒
	}

	// 显示下载信息
	fmt.Printf("下载链接: %s\n", *url)
	fmt.Printf("保存路径: %s\n", *output)
	fmt.Printf("文件大小: %.2f MB\n", float64(task.TotalSize)/1024/1024)
	fmt.Printf("下载线程: %d\n", task.SegmentCount)
	if *speedLimit > 0 {
		fmt.Printf("速度限制: %d KB/s\n", *speedLimit)
	}

	// 检查是否可以恢复之前的下载
	if task.canResume() {
		fmt.Println("检测到未完成的下载，正在恢复...")
		go task.monitorProgress()
		if err := task.resumeDownload(); err != nil {
			fmt.Printf("恢复下载失败: %v\n", err)
		}
	} else {
		// 开始新的下载
		go task.monitorProgress()
		if err := task.startDownload(); err != nil {
			fmt.Printf("下载失败: %v\n", err)
		}
	}
}
